package com.xjgz.aio.socket.transport;

import com.xjgz.aio.socket.MessageProcessor;
import com.xjgz.aio.socket.Protocol;
import com.xjgz.aio.socket.buffer.BufferPagePool;
import com.xjgz.aio.socket.enhance.EnhanceAsynchronousChannelProvider;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.SocketAddress;
import java.net.SocketOption;
import java.nio.channels.AsynchronousChannelGroup;
import java.nio.channels.AsynchronousSocketChannel;
import java.nio.channels.CompletionHandler;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

/**
 * AIO实现的客户端服务
 */
public class AioQuickClient {

    private static final ScheduledExecutorService CONNECT_TIMEOUT_EXECUTOR = Executors.newSingleThreadScheduledExecutor(r -> {
        Thread t = new Thread(r, "connection-timeout-monitor");
        t.setDaemon(true);
        return t;
    });

    /**
     * 网络连接的会话对象
     */
    private TcpAioSession session;

    /**
     * IO事件处理线程组
     */
    private AsynchronousChannelGroup asynchronousChannelGroup;

    /**
     * 绑定本地地址
     */
    private SocketAddress localAddress;

    /**
     * 连接超时时间
     */
    private int connectTimeout;

    /**
     * 客户端服务配置
     */
    private final IoServerConfig config = new IoServerConfig();

    /**
     * write内存池
     */
    private BufferPagePool writeBufferPool = null;

    /**
     * read内存池
     */
    private BufferPagePool readBufferPool = null;

    /**
     * 低内存模式
     */
    private boolean lowMemory = true;

    public <T> AioQuickClient(String host, int port, Protocol<T> protocol, MessageProcessor<T> processor) {
        config.setHost(host);
        config.setPort(port);
        config.setProtocol(protocol);
        config.setProcessor(processor);
    }

    public <A> void start(A attachment, CompletionHandler<AioSession, ? super A> handler) throws IOException {
        this.asynchronousChannelGroup =
                new EnhanceAsynchronousChannelProvider(lowMemory)
                        .openAsynchronousChannelGroup(2, Thread::new);
        start(asynchronousChannelGroup, attachment, handler);
    }

    public <A> void start(AsynchronousChannelGroup asynchronousChannelGroup, A attachment, CompletionHandler<AioSession, ? super A> handler) throws IOException {
        AsynchronousSocketChannel socketChannel = AsynchronousSocketChannel.open(asynchronousChannelGroup);
        if (connectTimeout > 0) {
            CONNECT_TIMEOUT_EXECUTOR.schedule(() -> {
                if (session == null) {
                    IOUtil.close(socketChannel);
                    shutdownNow();
                }
            }, connectTimeout, TimeUnit.MILLISECONDS);
        }
        if (writeBufferPool == null) {
            this.writeBufferPool = BufferPagePool.DEFAULT_BUFFER_PAGE_POOL;
        }
        if (readBufferPool == null) {
            this.readBufferPool = BufferPagePool.DEFAULT_BUFFER_PAGE_POOL;
        }

        if (config.getSocketOption() != null) {
            for (Map.Entry<SocketOption<Object>, Object> entry : config.getSocketOption().entrySet()) {
                socketChannel.setOption(entry.getKey(), entry.getValue());
            }
        }
        if (localAddress != null) {
            socketChannel.bind(localAddress);
        }

        socketChannel.connect(new InetSocketAddress(config.getHost(), config.getPort()), socketChannel, new CompletionHandler<>() {
            @Override
            public void completed(Void result, AsynchronousSocketChannel socketChannel) {
                try {
                    AsynchronousSocketChannel connectedChannel = null;
                    if (config.getMonitor() != null) {
                        connectedChannel = config.getMonitor().shouldAccept(socketChannel);
                    }
                    if (null == connectedChannel) {
                        throw new RuntimeException("NetMonitor refuse channel");
                    }
                    session = new TcpAioSession(connectedChannel, config, writeBufferPool, () -> readBufferPool.allocateSequentially(config.getReadBufferSize()));
                    handler.completed(session, attachment);
                } catch (Exception e) {
                    failed(e, socketChannel);
                }
            }

            @Override
            public void failed(Throwable exc, AsynchronousSocketChannel socketChannel) {
                try {
                    handler.failed(exc, attachment);
                } catch (Exception e) {
                    e.printStackTrace();
                } finally {
                    if (socketChannel != null) {
                        IOUtil.close(socketChannel);
                    }
                    shutdownNow();
                }
            }
        });
    }

    public AioSession start(AsynchronousChannelGroup asynchronousChannelGroup) throws IOException {
        CompletableFuture<AioSession> future = new CompletableFuture<>();
        start(asynchronousChannelGroup, future, new CompletionHandler<>() {

            @Override
            public void completed(AioSession session, CompletableFuture<AioSession> future) {
                if (future.isDone() || future.isCancelled()) {
                    session.close();
                } else {
                    future.complete(session);
                }
            }

            @Override
            public void failed(Throwable exc, CompletableFuture<AioSession> future) {
                future.completeExceptionally(exc);
            }
        });
        try {
            return future.get();
        } catch (Exception e) {
            future.cancel(false);
            shutdownNow();
            throw new IOException(e.getCause() == null ? e : e.getCause());
        }
    }

    public TcpAioSession getSession() {
        return session;
    }

    public AioSession start() throws IOException {
        this.asynchronousChannelGroup =
                new EnhanceAsynchronousChannelProvider(lowMemory)
                        .openAsynchronousChannelGroup(2, Thread::new);
        return start(this.asynchronousChannelGroup);
    }

    public void shutdown() {
        shutdown0(false);
    }

    public void shutdownNow() {
        shutdown0(true);
    }

    private void shutdown0(boolean immediate) {
        if (session != null) {
            session.close(immediate);
            session = null;
        }
        // 仅client内部创建的ChannelGroup需要shutdown
        synchronized (this) {
            if (asynchronousChannelGroup != null) {
                asynchronousChannelGroup.shutdown();
                asynchronousChannelGroup = null;
            }
        }
    }

    /**
     * 设置读缓冲区大小
     *
     * @param size 单位字节
     */
    public AioQuickClient setReadBufferSize(int size) {
        this.config.setReadBufferSize(size);
        return this;
    }

    /**
     * 设置Socket的TCP参数配置
     *
     * @param socketOption StandardSocketOptions中的选项
     * @param value        选项值
     */
    public AioQuickClient setOption(SocketOption<Object> socketOption, Object value) {
        config.setSocketOption(socketOption, value);
        return this;
    }

    /**
     * 绑定本地地址
     *
     * @param local 若为null则由系统自动获取
     * @param port  若为0则由系统指定
     */
    public AioQuickClient bindLocal(String local, int port) {
        this.localAddress = local == null ? new InetSocketAddress(port) : new InetSocketAddress(local, port);
        return this;
    }

    /**
     * 设置内存池
     * 通过该方法设置的内存池，在AioQuickClient执行shutdown时不会触发内存池的释放
     */
    public AioQuickClient setBufferPagePool(BufferPagePool bufferPagePool) {
        return setBufferPagePool(bufferPagePool, bufferPagePool);
    }

    public AioQuickClient setBufferPagePool(BufferPagePool readBufferPool, BufferPagePool writeBufferPool) {
        this.writeBufferPool = readBufferPool;
        this.readBufferPool = writeBufferPool;
        return this;
    }

    /**
     * 设置输出缓冲区容量
     *
     * @param chunkSize  单个内存块的大小
     * @param chunkCount 内存块数量
     */
    public AioQuickClient setWriteBuffer(int chunkSize, int chunkCount) {
        config.setWriteChunkSize(chunkSize);
        config.setWriteChunkCount(chunkCount);
        return this;
    }

    public AioQuickClient connectTimeout(int timeout) {
        this.connectTimeout = timeout;
        return this;
    }

    public AioQuickClient disableLowMemory() {
        this.lowMemory = false;
        return this;
    }
}
